import librosa
import numpy as np
import keras
from keras.models import Sequential
from keras.layers import Dense
import math


print('读取训练信号\n')
for x in range(500):
    y, sr = librosa.load(r'C:\Users\WangCan\PycharmProjects\speech\venv\datasource\s1\S1_'+str(x+1)+'.wav', sr=None)
    y = y.reshape(1, len(y))
    if(x == 0):
        maleSignal = y
    else:
        maleSignal = np.c_[maleSignal, y]
    y, sr = librosa.load(r'C:\Users\WangCan\PycharmProjects\speech\venv\datasource\s16\S16_' + str(x + 1) + '.wav', sr=None)
    y = y.reshape(1, len(y))
    if (x == 0):
        femaleSignal = y
    else:
        femaleSignal = np.c_[femaleSignal, y]
'#18608750 21256500'
print(len(maleSignal[0]), len(femaleSignal[0]))
maleSignal = np.array(maleSignal[0][0:18432000]).reshape(1, -1)
femaleSignal = np.array(femaleSignal[0][0:18432000]).reshape(1, -1)
print(len(maleSignal[0]))
mixedSignal = np.array(maleSignal[0]+femaleSignal[0]).reshape(1, -1)

print("读取测试信号\n")
for x in range(10):
    y, sr = librosa.load(r'C:\Users\WangCan\PycharmProjects\speech\venv\datasource\s1\S1_'+str(x+501)+'.wav', sr=None)
    y = y.reshape(1, len(y))
    if(x == 0):
        cMaleSignal = y
    else:
        cMaleSignal = np.c_[cMaleSignal, y]
    y, sr = librosa.load(r'C:\Users\WangCan\PycharmProjects\speech\venv\datasource\s16\S16_' + str(x + 501) + '.wav',sr=None)
    y = y.reshape(1, len(y))
    if (x == 0):
        cFemaleSignal = y
    else:
        cFemaleSignal = np.c_[cFemaleSignal, y]
'#364010 427010'
print(len(cMaleSignal[0]), len(cFemaleSignal[0]))
cMaleSignal = np.array(cMaleSignal[0][0:363520]).reshape(1, -1)
cFemaleSignal = np.array(cFemaleSignal[0][0:363520]).reshape(1, -1)
print(len(cMaleSignal[0]), len(cMaleSignal[0]) == len(cFemaleSignal[0]))
cMixedSignal = np.array(cMaleSignal[0]+cFemaleSignal[0]).reshape(1, -1)

print("信号变换：")
#混合信号
bMixedSignal = librosa.stft(mixedSignal[0], n_fft=512)
bMixedSignalAmplitude = np.abs(bMixedSignal)
TrainMixedSignal = bMixedSignalAmplitude.transpose()
# 男声信号
bMaleSignal = librosa.stft(maleSignal[0], n_fft=512)
bMaleSignalAmplitude = np.abs(bMaleSignal)
TrainMaleSignal = bMaleSignalAmplitude.transpose()
#测试混合信号
tMixedMaleSignal = librosa.stft(cMaleSignal[0], n_fft=512)
tMixedSignal = librosa.stft(cMixedSignal[0], n_fft=512)
tMixedSignalAmplitude = np.abs(tMixedSignal)
tMixedSignalAngle = np.angle(tMixedMaleSignal)
# tMixedSignalAngle = np.angle(tMixedSignal)
TestMixedSignal = tMixedSignalAmplitude.transpose()
print("信号变换完成")

print("开始训练：")
X_train = TrainMixedSignal
Y_train = TrainMaleSignal
X_test = TestMixedSignal
model = Sequential()
model.add(Dense(1024, activation='relu', input_dim=257))
model.add(Dense(1024, activation='selu'))
model.add(Dense(1024, activation='sigmoid'))
model.add(Dense(257, activation='relu'))
# sgd = keras.optimizers.SGD(lr=0.001, momentum=0.0, decay=0.0, nesterov=False)
# ADAM = keras.optimizers.Adam(lr=0.001, beta_1=0.9, beta_2=0.999, epsilon=None, decay=0.0, amsgrad=False)
# NADAM =keras.optimizers.Nadam(lr=0.002, beta_1=0.9, beta_2=0.999)
# ADAMAX=keras.optimizers.Adamax(lr=0.002, beta_1=0.9, beta_2=0.999)
# adaBound = AdaBound(model.parameters(), lr=1e-3, final_lr=0.1)
# Adadelta=keras.optimizers.Adadelta(lr=1.0, rho=0.95)
model.compile(optimizer='adam', loss='mse', metrics=['mae'])
model.fit(X_train, Y_train, epochs=200, batch_size=256, verbose=2)
Y_pred = model.predict(X_test, batch_size=256)
print('训练完成\n')

print("语音输出开始：")
Y_pred = Y_pred.transpose()
rMaleSignal = Y_pred*np.power(math.e, 1j*tMixedSignalAngle)
rebuildMaleSignal = librosa.istft(rMaleSignal)
librosa.output.write_wav(r".\speech\mapping\rebuild_male.wav", rebuildMaleSignal, sr=25000)
librosa.output.write_wav(r".\speech\mapping\original_mixed.wav", cMixedSignal[0], sr=25000)
print("语音输出结束；")